{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Tutorial Part 5: Creating Models with TensorFlow and PyTorch\n",
    "\n",
    "In the tutorials so far, we have used standard models provided by DeepChem.  This is fine for many applications, but sooner or later you will want to create an entirely new model with an architecture you define yourself.  DeepChem provides integration with both TensorFlow (Keras) and PyTorch, so you can use it with models from either of these frameworks.\n",
    "\n",
    "## Colab\n",
    "\n",
    "This tutorial and the rest in this sequence are designed to be done in Google colab. If you'd like to open this notebook in colab, you can use the following link.\n",
    "\n",
    "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/deepchem/deepchem/blob/master/examples/tutorials/05_Creating_Models_with_TensorFlow_and_PyTorch.ipynb)\n",
    "\n",
    "## Setup\n",
    "\n",
    "To run DeepChem within Colab, you'll need to run the following installation commands. This will take about 5 minutes to run to completion and install your environment. You can of course run this tutorial locally if you prefer. In that case, don't run these cells since they will download and install Anaconda on your local machine."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!curl -Lo conda_installer.py https://raw.githubusercontent.com/deepchem/deepchem/master/scripts/colab_install.py\n",
    "import conda_installer\n",
    "conda_installer.install()\n",
    "!/root/miniconda/bin/conda info -e"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install --pre deepchem"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "There are actually two different approaches you can take to using TensorFlow or PyTorch models with DeepChem.  It depends on whether you want to use TensorFlow/PyTorch APIs or DeepChem APIs for training and evaluating your model.  For the former case, DeepChem's `Dataset` class has methods for easily adapting it to use with other frameworks.  `make_tf_dataset()` returns a `tensorflow.data.Dataset` object that iterates over the data.  `make_pytorch_dataset()` returns a `torch.utils.data.IterableDataset` that iterates over the data.  This lets you use DeepChem's datasets, loaders, featurizers, transformers, splitters, etc. and easily integrate them into your existing TensorFlow or PyTorch code.\n",
    "\n",
    "But DeepChem also provides many other useful features.  The other approach, which lets you use those features, is to wrap your model in a DeepChem `Model` object.  Let's look at how to do that.\n",
    "\n",
    "## KerasModel\n",
    "\n",
    "`KerasModel` is a subclass of DeepChem's `Model` class.  It acts as a wrapper around a `tensorflow.keras.Model`.  Let's see an example of using it.  For this example, we create a simple sequential model consisting of two dense layers."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import deepchem as dc\n",
    "import tensorflow as tf\n",
    "\n",
    "keras_model = tf.keras.Sequential([\n",
    "    tf.keras.layers.Dense(1000, activation='relu'),\n",
    "    tf.keras.layers.Dropout(rate=0.5),\n",
    "    tf.keras.layers.Dense(1)\n",
    "])\n",
    "model = dc.models.KerasModel(keras_model, dc.models.losses.L2Loss())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For this example, we used the Keras `Sequential` class.  Our model consists of a dense layer with ReLU activation, 50% dropout to provide regularization, and a final layer that produces a scalar output.  We also need to specify the loss function to use when training the model, in this case L<sub>2</sub> loss.  We can now train and evaluate the model exactly as we would with any other DeepChem model.  For example, let's load the Delaney solubility dataset.  How does our model do at predicting the solubilities of molecules based on their extended-connectivity fingerprints (ECFPs)?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "training set score: {'pearson_r2_score': 0.9787445597470444}\n",
      "test set score: {'pearson_r2_score': 0.736905850092889}\n"
     ]
    }
   ],
   "source": [
    "tasks, datasets, transformers = dc.molnet.load_delaney(featurizer='ECFP', splitter='random')\n",
    "train_dataset, valid_dataset, test_dataset = datasets\n",
    "model.fit(train_dataset, nb_epoch=50)\n",
    "metric = dc.metrics.Metric(dc.metrics.pearson_r2_score)\n",
    "print('training set score:', model.evaluate(train_dataset, [metric]))\n",
    "print('test set score:', model.evaluate(test_dataset, [metric]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## TorchModel\n",
    "\n",
    "`TorchModel` works just like `KerasModel`, except it wraps a `torch.nn.Module`.  Let's use PyTorch to create another model just like the previous one and train it on the same data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "training set score: {'pearson_r2_score': 0.9798256761766225}\n",
      "test set score: {'pearson_r2_score': 0.7256745385608444}\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "\n",
    "pytorch_model = torch.nn.Sequential(\n",
    "    torch.nn.Linear(1024, 1000),\n",
    "    torch.nn.ReLU(),\n",
    "    torch.nn.Dropout(0.5),\n",
    "    torch.nn.Linear(1000, 1)\n",
    ")\n",
    "model = dc.models.TorchModel(pytorch_model, dc.models.losses.L2Loss())\n",
    "\n",
    "model.fit(train_dataset, nb_epoch=50)\n",
    "print('training set score:', model.evaluate(train_dataset, [metric]))\n",
    "print('test set score:', model.evaluate(test_dataset, [metric]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Computing Losses\n",
    "\n",
    "Now let's see a more advanced example.  In the above models, the loss was computed directly from the model's output.  Often that is fine, but not always.  Consider a classification model that outputs a probability distribution.  While it is possible to compute the loss from the probabilities, it is more numerically stable to compute it from the logits.\n",
    "\n",
    "To do this, we create a model that returns multiple outputs, both probabilities and logits.  `KerasModel` and `TorchModel` let you specify a list of \"output types\".  If a particular output has type `'prediction'`, that means it is a normal output that should be returned when you call `predict()`.  If it has type `'loss'`, that means it should be passed to the loss function in place of the normal outputs.\n",
    "\n",
    "Sequential models do not allow multiple outputs, so instead we use a subclassing style model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "class ClassificationModel(tf.keras.Model):\n",
    "    \n",
    "    def __init__(self):\n",
    "        super(ClassificationModel, self).__init__()\n",
    "        self.dense1 = tf.keras.layers.Dense(1000, activation='relu')\n",
    "        self.dense2 = tf.keras.layers.Dense(1)\n",
    "\n",
    "    def call(self, inputs, training=False):\n",
    "        y = self.dense1(inputs)\n",
    "        if training:\n",
    "            y = tf.nn.dropout(y, 0.5)\n",
    "        logits = self.dense2(y)\n",
    "        output = tf.nn.sigmoid(logits)\n",
    "        return output, logits\n",
    "\n",
    "keras_model = ClassificationModel()\n",
    "output_types = ['prediction', 'loss']\n",
    "model = dc.models.KerasModel(keras_model, dc.models.losses.SigmoidCrossEntropy(), output_types=output_types)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can train our model on the BACE dataset.  This is a binary classification task that tries to predict whether a molecule will inhibit the enzyme BACE-1."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "training set score: {'roc_auc_score': 0.9996116504854369}\n",
      "test set score: {'roc_auc_score': 0.7701992753623188}\n"
     ]
    }
   ],
   "source": [
    "tasks, datasets, transformers = dc.molnet.load_bace_classification(feturizer='ECFP', split='scaffold')\n",
    "train_dataset, valid_dataset, test_dataset = datasets\n",
    "model.fit(train_dataset, nb_epoch=100)\n",
    "metric = dc.metrics.Metric(dc.metrics.roc_auc_score)\n",
    "print('training set score:', model.evaluate(train_dataset, [metric]))\n",
    "print('test set score:', model.evaluate(test_dataset, [metric]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Other Features\n",
    "\n",
    "`KerasModel` and `TorchModel` have lots of other features.  Here are some of the more important ones.\n",
    "\n",
    "- Automatically saving checkpoints during training.\n",
    "- Logging progress to the console, to [TensorBoard](https://www.tensorflow.org/tensorboard), or to [Weights & Biases](https://docs.wandb.com/).\n",
    "- Custom loss functions that you define with a function of the form `f(outputs, labels, weights)`.\n",
    "- Early stopping using the `ValidationCallback` class.\n",
    "- Loading parameters from pre-trained models.\n",
    "- Estimating uncertainty in model outputs.\n",
    "- Identifying important features through saliency mapping.\n",
    "\n",
    "By wrapping your own models in a `KerasModel` or `TorchModel`, you get immediate access to all these features.  See the API documentation for full details on them.\n",
    "\n",
    "# Congratulations! Time to join the Community!\n",
    "\n",
    "Congratulations on completing this tutorial notebook! If you enjoyed working through the tutorial, and want to continue working with DeepChem, we encourage you to finish the rest of the tutorials in this series. You can also help the DeepChem community in the following ways:\n",
    "\n",
    "## Star DeepChem on [GitHub](https://github.com/deepchem/deepchem)\n",
    "This helps build awareness of the DeepChem project and the tools for open source drug discovery that we're trying to build.\n",
    "\n",
    "## Join the DeepChem Gitter\n",
    "The DeepChem [Gitter](https://gitter.im/deepchem/Lobby) hosts a number of scientists, developers, and enthusiasts interested in deep learning for the life sciences. Join the conversation!"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
